from __future__ import division
import numpy as np
import pylab as pl
import cv2.cv as cv

class snake_energy:
	def __init__(self, snake, image,w=(1,1,1)):
		''' Class that calculates the energy of a contour in position
			*snake* in *image*. The energy is stored in class variable
			E. The class is in very early trial stage.'''
		self.snake = np.asarray(snake)
		mat = cv.CreateMat(image.height,image.width,cv.CV_32FC1 )
		self.w = w
		cv.Convert( image, mat )
		self.image = np.asarray(mat)

		self.N = self.snake.shape[0]
		self.E = self.internal_energy() + self.w[2]*self.image_energy()
		
	def internal_energy(self):
		### The internal bending energy of the snake
		dVds = np.diff(self.snake,axis=0)
		dist = np.sqrt(dVds[:,0]**2 + dVds[:,1]**2)
		d2Vds2 = np.diff(dVds,axis=0)
		dist2 = np.sqrt(d2Vds2[:,0]**2 + d2Vds2[:,1]**2)
		self.internal1 = np.sum(dist)*(1.0/(self.N))
		self.internal2 = np.sum(dist2)*(1.0/(self.N-1))

		return self.w[0]*np.sum(dist)*(1.0/(self.N))+self.w[1]*np.sum(dist2)*(1.0/(self.N-1))
		
	def image_energy(self):
		### Image energy for the snake
		Eline = self.image[self.snake[:,0],self.snake[:,1]]
		grad = np.gradient(self.image)
		grad = np.max([abs(grad[0]),abs(grad[1])], axis=0)
		Eedge = grad[self.snake[:,0],self.snake[:,1]]

		return np.sum(Eline)*(1.0/(self.N+1)) - np.sum(Eedge)*(1.0/(self.N+1))
